# 本脚本用于针对 关联矩阵的异常检测


import numpy as np
import pandas as pd

from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
from keras.models import Model
from keras.models import load_model
from keras import optimizers

import re
import matplotlib.pyplot as plt




def getMatrix(x_train):
    w = 10
    n_samples = x_train.shape[0] # 8000
    x_rs = []
    for k in range(n_samples):
        tmp = x_train[k] # 14, 10
        m = np.corrcoef(tmp) # 14, 14
        x_rs.append(m)

    x_rs = np.array(x_rs)
    x_rs = np.reshape(x_rs, (n_samples, x_rs.shape[1], x_rs.shape[2], 1))
    return x_rs


def train(raw_data):
    print('train...start')
    raw_data = raw_data[:, 0:8000]
    train_len = raw_data.shape[1]
    x_train = []
    for i in range(0, train_len, 10):
        x_train.append(raw_data[:, i:i+10])

    x_train = np.array(x_train)

    x_train = getMatrix(x_train) # 8000,14,10,1 -> 8000,14,14,1

    #
    input_img = Input(shape=(14, 14, 1))
    # Model Construction
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(input_img)
    x = MaxPooling2D((1, 1), padding='same')(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    encoded = MaxPooling2D((1, 1), padding='same')(x)
    # At this point the representation is (14, 14, 32)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(encoded)
    x = UpSampling2D((1, 1))(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    x = UpSampling2D((1, 1))(x)
    decoded = Conv2D(1, (3, 3), activation='relu', padding='same')(x)
    autoencoder = Model(input_img, decoded)

    autoencoder.compile(optimizer='adam', loss='mse')

    autoencoder.fit(x_train, x_train, epochs=100, batch_size=256, shuffle=True, validation_split=0.05)

    autoencoder.save('../Auto_Model/model_AEX.h5')

    print('train...end')


def test(raw_data):
    print('test...start')

    model = load_model('../Auto_Model/model_AEX.h5')
    raw_data = raw_data[:, 10000:]
    test_len = raw_data.shape[1]

    x_test = []
    for i in range(0, test_len, 10):
        x_test.append(raw_data[:, i:i + 10])
    x_test = np.array(x_test)

    x_test = getMatrix(x_test)

    predicted_value = model.predict(x_test)

    error = np.abs(predicted_value - x_test)  #

    # np.save('../rs/AE-M_error.npy', error)

    test_anomaly_score = []

    n_samples = error.shape[0]

    for i in range(n_samples):
        tmp = np.reshape(error[i], (14, 14))
        error_mean = np.mean(tmp)
        test_anomaly_score.append(error_mean)

    test_anomaly_score = np.array(test_anomaly_score)


    anomaly_pos = np.zeros(4) #针对TestSet1-1
    anomaly_span = [90]
    root_cause_f = open("../data/test_anomaly.csv", "r")

    row_index = 0

    for line in root_cause_f:
        line = line.strip()
        anomaly_axis = int(re.split(',', line)[0])
        anomaly_pos[row_index] = anomaly_axis / 10 - 1000 - anomaly_span[row_index % 1] / 10
        row_index += 1
    root_cause_f.close()

    fig, axes = plt.subplots()

    test_num = 1000
    plt.xticks(fontsize=25)

    np.savetxt('../rs/AE-x.csv', test_anomaly_score, delimiter=',')

    plt.plot(test_anomaly_score, 'black', linewidth=2)
    threshold = np.full((test_num), 0.01)
    axes.plot(threshold, color='black', linestyle='--', linewidth=2)
    for k in range(len(anomaly_pos)):
        axes.axvspan(anomaly_pos[k], anomaly_pos[k] + anomaly_span[k % 1] / 10, color='red',
                    linewidth=2)  # anomaly为起始位置，anomaly_pos[k]+anomaly_span[k%3]/gap_time为终止位置

    plt.xlabel('Test Time', fontsize=25)
    plt.ylabel('Anomaly Score', fontsize=25)
    axes.spines['right'].set_visible(False)
    axes.spines['top'].set_visible(False)
    axes.yaxis.set_ticks_position('left')
    axes.xaxis.set_ticks_position('bottom')
    fig.subplots_adjust(bottom=0.25)
    fig.subplots_adjust(left=0.25)
    plt.title("AutoEncoder-Coef", size=25)
    # plt.show()
    plt.savefig('../rs/AE-x.png', dpi=600)
    print('test...end')


if __name__ == '__main__':
    raw_data = pd.read_csv('../data/data_testAno.csv', header=None).values # 针对TestSet1-1
    # raw_data = pd.read_csv('../data/data_testAno_2.csv', header=None).values  # 针对TestSet1-2

    # 数据重组
    raw_data = raw_data[(0,1,2,3,4,5,6,7,9,10,11,12,13,14),]

    # min-max
    max_value = np.max(raw_data, axis=1)
    min_value = np.min(raw_data, axis=1)
    raw_data = (np.transpose(raw_data) - min_value) / (max_value - min_value + 1e-6)
    raw_data = np.transpose(raw_data)

    train(raw_data)
    test(raw_data)